package com.duyus.ai.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * MD文档加载器
 */
@Component
@Slf4j
public class MdDocumentLoader {

    private final ResourcePatternResolver resourcePatternResolver;

    // 依赖注入
    public MdDocumentLoader(ResourcePatternResolver resourcePatternResolver) {
        this.resourcePatternResolver = resourcePatternResolver;
    }

    /**
     * 加载多篇 Markdown 文档
     * 通过 Mar⁠kdownDocumentReaderConfig 文‌档加载配置来指定读取文档的细节，比如是否读取代码块、引用块等。
     * 特别需要注意的是，我们还指定了额外的元信息配置，提取文档的文件名（fileName）作为文档的元信息‌，可以便于后续知识库实现更精确的检索。
     * @return
     */
    public List<Document> loadMarkdowns() {
        List<Document> allDocuments = new ArrayList<>();
        try {
            // 从 classpath 下所有 jar 包或目录中，匹配 docs 文件夹下的 .md 文件
            Resource[] resources = resourcePatternResolver.getResources("classpath:document/*.md");
            for (Resource resource : resources) {
                String filename = resource.getFilename();
                // 提取文档倒数第 3 和第 2 个字作为标签
                String status = filename.substring(filename.length() - 6, filename.length() - 4);
                //  解释 后面代码作用  MarkdownDocumentReaderConfig 是 Markdown 文档的读取配置
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        // 是否根据分割线 创建文档 默认是根据标题分割 又加入了 分割线
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(false)
                        .withIncludeBlockquote(false)
                        .withAdditionalMetadata("filename", filename)
                        .withAdditionalMetadata("status", status)
                        .build();
                MarkdownDocumentReader markdownDocumentReader = new MarkdownDocumentReader(resource, config);
                // MarkdownDocumentReader 分割文档，返回的是一个流
                allDocuments.addAll(markdownDocumentReader.get());
            }
        } catch (IOException e) {
           log.error("Markdown 文档加载失败", e);
        }
        return allDocuments;
    }
}
